#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <iostream>
#include <cstring>
#include <cmath>
#include <stdio.h>
#include <omp.h>

// Template structure for 3D point representation
// It supports basic operations like initialization and printing
template <typename T>
struct Point
{
    T x, y, z; // Coordinates of the point in 3D space

    // Default constructor initializing point to origin
    Point() : x(0), y(0), z(0) {}

    // Parameterized constructor for initializing point to given coordinates
    Point(T x, T y, T z) : x(x), y(y), z(z) {}

    // Copy constructor for duplicating a Point object
    Point(const Point& p) : x(p.x), y(p.y), z(p.z) {}

    // Method to print point coordinates
    void print() const
    {
        std::cout << "(" << x << ", " << y << ", " << z << ")" << std::endl;
    }
};

// Function to calculate the squared distance between two 3D points
// Using squared distance avoids unnecessary computation of square root
template <typename T>
inline T distance2(const Point<T>& p1, const Point<T>& p2)
{
    T dist2 = (p1.x - p2.x) * (p1.x - p2.x) + (p1.y - p2.y) * (p1.y - p2.y) + (p1.z - p2.z) * (p1.z - p2.z);
    return dist2;
}

// Function to read the grid size from a file (SPECGRID keyword section)
// NX, NY, NZ represent the dimensions of the grid
template <typename T>
void readCornerPointGridSize(const std::string filename, int& NX, int& NY, int& NZ)
{
    std::ifstream file(filename);
    if (!file.is_open()) // Check if the file can be opened
    {
        std::cerr << "Error: could not open file " << filename << std::endl;
        return;
    }
    std::string line;
    while (std::getline(file, line)) // Iterate through each line of the file
    {
        if (line.find("SPECGRID") != std::string::npos) // Find the SPECGRID section
        {
            std::getline(file, line); // Read the next line containing grid dimensions
            std::istringstream iss(line);
            iss >> NX >> NY >> NZ; // Extract grid dimensions
            break;
        }
    }
    file.close(); // Close the file after reading
}

// Function to read corner point grid data from a file
// This includes reading coordinates, ZCORN values, and activity (ACTNUM) indicators
// NX, NY, NZ are the grid dimensions
template <typename T>
void readCornerPointGrid(const std::string filename, const int NX, const int NY, const int NZ, \
std::vector<T>& coords, std::vector<T>& zcorns, std::vector<bool>& acts)
{
    std::ifstream file(filename);
    if (!file.is_open()) // Check if the file can be opened
    {
        std::cerr << "Error: could not open file " << filename << std::endl;
        return;
    }
    std::string line;

    T x, y, z; // Temporary variables to hold coordinate values

    // Read COORD data (X, Y coordinates of grid corners)
    while (std::getline(file, line))
    {
        if (line.find("COORD") != std::string::npos)
        {
            for (int i = 0; i < (NX + 1) * (NY + 1); i++)
            {
                std::getline(file, line);
                std::istringstream iss(line);
                iss >> x >> y;
                coords[i * 2] = x;   // Store X coordinate
                coords[i * 2 + 1] = y; // Store Y coordinate
            }
            break;
        }
    }
    
    // Read ZCORN data (Z coordinates of grid corners)
    while (std::getline(file, line))
    {
        if (line.find("ZCORN") != std::string::npos)
        {
            std::getline(file, line, '/'); // Extract data up to delimiter '/'
            std::istringstream iss(line);
            int cnt = 0;
            while (iss >> z)
            {
                zcorns[cnt++] = z; // Store Z coordinate
            }
            break;
        }   
    }

    // Read ACTNUM data (grid cell activity indicators)
    while (std::getline(file, line))
    {
        if (line.find("ACTNUM") != std::string::npos)
        {
            std::getline(file, line, '/'); // Extract data up to delimiter '/'
            std::istringstream iss(line);
            std::string token;
            int actnum;
            bool act;
            int cnt = 0;
            while (iss >> token)
            {
                auto pos = token.find('*');
                if (pos != std::string::npos)
                {
                    actnum = std::stoi(token.substr(0, pos)); // Number of repetitions
                    act = std::stoi(token.substr(pos + 1));   // Active/inactive flag
                    for (int i = 0; i < actnum; i++)
                    {
                        acts[cnt++] = act;
                    }
                }
            }
            break;
        }
    }

    file.close(); // Close the file after reading
}

// Function to assemble 3D points for the grid from coordinate and ZCORN data
// Also updates isPointExist array to mark valid points
// Uses OpenMP to parallelize point processing
template <typename T>
void assemblePoints(const std::vector<T>& coords, const std::vector<T>& zcorns, const std::vector<bool> acts, \
int NX, int NY, int NZ, std::vector<Point<T>>& points, std::vector<bool>& isPointExist)
{   
    #pragma omp parallel for
    for (int k = 0; k < NZ; k++)
    {
        for (int j = 0; j < NY; j++)
        {
            for (int i = 0; i < NX; i++)
            {
                int cellIndex = k * NY * NX + j * NX + i; // Compute cell index
                if (acts[cellIndex] == false) continue; // Skip inactive cells
                for (int index = 0; index < 8; index++) // Process 8 corners of the cell
                {
                    int x = i + (index & 1); // Determine X offset
                    int y = j + ((index & 2) >> 1); // Determine Y offset
                    int z = k + ((index & 4) >> 2); // Determine Z offset
                    int idx = z * (NY + 1) * (NX + 1) + y * (NX + 1) + x; // Global point index
                    if (isPointExist[idx] == false)
                    {
                        T px = coords[y * (NX + 1) * 2 + x * 2 + 0]; // Fetch X coordinate
                        T py = coords[y * (NX + 1) * 2 + x * 2 + 1]; // Fetch Y coordinate
                        T pz = zcorns[k * NX * NY * 8 + ((index & 4) >> 2) * NX * NY * 4 + j * NX * 4 + ((index & 2) >> 1) * NX * 2 + i * 2 + (index & 1)];
                        points[idx] = Point<T>(px, py, pz); // Create Point object
                        isPointExist[idx] = true; // Mark point as valid
                    }
                }
            }
        }
    }
}

template <typename T>
void findFault(const std::vector<T>& coords, const std::vector<bool>& isPointExist, int NX, int NY, int NZ, std::vector<bool>& faults)
{
    T fault_data[4]; // Array to store coordinates of potential fault points
    T dx = coords[2] - coords[0]; // Calculate grid spacing in x-direction
    T dy = coords[(NX + 1) * 2 + 1] - coords[1]; // Calculate grid spacing in y-direction
    T x0 = coords[0]; // Starting x-coordinate
    T y0 = coords[1]; // Starting y-coordinate
    T eps = sqrt(dx * dx + dy * dy) * 1e-2; // Tolerance for detecting faults, based on grid spacing
    int cnt = 0; // Counter to track fault points

    // Loop through grid points to identify faults
    for (int j = 0; j < NY + 1; j++)
    {
        for (int i = 0; i < NX + 1; i++)
        {
            int idx = j * (NX + 1) + i; // Compute linear index for current grid point
            if (!isPointExist[idx]) continue; // Skip points that do not exist

            // Extract current x and y coordinates
            T x = coords[j * (NX + 1) * 2 + i * 2 + 0];
            T y = coords[j * (NX + 1) * 2 + i * 2 + 1];

            // Check if the point deviates significantly from the regular grid spacing
            if (fabs(x - x0 - i * dx) > eps || fabs(y - y0 - j * dy) > eps)
            {
                fault_data[cnt++] = x; // Store x-coordinate of the fault point
                fault_data[cnt++] = y; // Store y-coordinate of the fault point
                if (cnt == 4) break; // Stop if we have identified 4 fault points
            }
        }
        if (cnt == 4) break;
    }

    // Calculate the fault line equation y = kx + b
    T k = (fault_data[3] - fault_data[1]) / (fault_data[2] - fault_data[0]); // Slope
    T b = fault_data[1] - k * fault_data[0]; // Intercept

    // Parallel loop to mark faults in the faults array
    #pragma omp parallel for
    for (int j = 0; j < NY + 1; j++)
    {
        for (int i = 0; i < NX + 1; i++)
        {
            int idx = j * (NX + 1) + i; // Compute linear index for current grid point
            if (!isPointExist[idx]) continue; // Skip points that do not exist

            // Extract current x and y coordinates
            T x = coords[j * (NX + 1) * 2 + i * 2 + 0];
            T y = coords[j * (NX + 1) * 2 + i * 2 + 1];

            // Calculate expected y-coordinate based on the fault line equation
            T y1 = k * x + b;

            // Mark point as part of a fault if its y-coordinate is close to the fault line
            if (fabs(y - y1) < eps)
            {
                faults[j * (NX + 1) + i] = true;
            }
        }
    }
}

// Function to convert a hexahedron into tetrahedrons based on distances between points
// - Template function to support various data types for point coordinates
// - Input:
//   * points_index: Array of point indices (length = 8, representing corners of a hexahedron)
//   * points_coords: Array of Point<T> objects, representing coordinates of the points
//   * tet_par: Output array for tetrahedron vertex indices (length = 24, 6 tetrahedrons x 4 vertices each)
template <typename T>
inline void hexahedron26tetrahedron(const int* points_index, const Point<T>* points_coords, int* tet_par)
{
    // Compare squared distances between specific points to decide the split strategy
    if (distance2(points_coords[1], points_coords[2]) >= distance2(points_coords[0], points_coords[3]))
    {
        // First splitting strategy: Assign indices for 6 tetrahedrons
        tet_par[0] = points_index[2];
        tet_par[1] = points_index[3];
        tet_par[2] = points_index[0];
        tet_par[3] = points_index[4];

        tet_par[4] = points_index[2];
        tet_par[5] = points_index[4];
        tet_par[6] = points_index[6];
        tet_par[7] = points_index[3];

        tet_par[8] = points_index[3];
        tet_par[9] = points_index[4];
        tet_par[10] = points_index[6];
        tet_par[11] = points_index[7];

        tet_par[12] = points_index[3];
        tet_par[13] = points_index[5];
        tet_par[14] = points_index[4];
        tet_par[15] = points_index[7];

        tet_par[16] = points_index[1];
        tet_par[17] = points_index[3];
        tet_par[18] = points_index[5];
        tet_par[19] = points_index[4];

        tet_par[20] = points_index[1];
        tet_par[21] = points_index[0];
        tet_par[22] = points_index[3];
        tet_par[23] = points_index[4];
    }
    else
    {
        // Alternative splitting strategy: Assign indices for 6 tetrahedrons
        tet_par[0] = points_index[2];
        tet_par[1] = points_index[1];
        tet_par[2] = points_index[0];
        tet_par[3] = points_index[6];

        tet_par[4] = points_index[0];
        tet_par[5] = points_index[4];
        tet_par[6] = points_index[6];
        tet_par[7] = points_index[1];

        tet_par[8] = points_index[1];
        tet_par[9] = points_index[4];
        tet_par[10] = points_index[6];
        tet_par[11] = points_index[5];

        tet_par[12] = points_index[1];
        tet_par[13] = points_index[5];
        tet_par[14] = points_index[6];
        tet_par[15] = points_index[7];

        tet_par[16] = points_index[1];
        tet_par[17] = points_index[3];
        tet_par[18] = points_index[7];
        tet_par[19] = points_index[6];

        tet_par[20] = points_index[1];
        tet_par[21] = points_index[2];
        tet_par[22] = points_index[3];
        tet_par[23] = points_index[6];
    }
}

// Function to generate a VTK file from grid data
// - Outputs a visualization of the grid structure as tetrahedrons
// - Input:
//   * filename: Output file path
//   * points: Coordinates of all points in the grid
//   * acts: Boolean array indicating active grid cells
//   * NX, NY, NZ: Number of cells along X, Y, Z directions
//   * isPointExist: Boolean array indicating existing points in the grid
//   * faults: Boolean array representing fault information
template <typename T>
void print2vtk(const std::string filename, const std::vector<Point<T>>& points, const std::vector<bool> acts, \
const int NX, const int NY, const int NZ, const std::vector<bool>& isPointExist, const std::vector<bool>& faults)
{
    // Open output file
    FILE* file = fopen(filename.c_str(), "w");
    if (file == NULL)
    {
        std::cerr << "Error: could not open file " << filename << std::endl;
        return;
    }

    // Write VTK file header
    fprintf(file, "# vtk DataFile Version 3.0\n");
    fprintf(file, "Grid_E in vtk\n");
    fprintf(file, "ASCII\n");
    fprintf(file, "DATASET UNSTRUCTURED_GRID\n\n");

    // Assign indices to points and count valid points
    std::vector<int> points_index((NX + 1) * (NY + 1) * (NZ + 1), 0);
    int point_num = 0;
    for (int i = 0; i < (NX + 1) * (NY + 1) * (NZ + 1); i++)
    {
        if (isPointExist[i])
        {
            point_num++;
        }
    }
    fprintf(file, "POINTS %d float\n", point_num);

    // Output valid points and build mapping from grid index to VTK point index
    point_num = 0;
    for (int k = 0; k < NZ + 1; k++)
    {
        for (int j = 0; j < NY + 1; j++)
        {
            for (int i = 0; i < NX + 1; i++)
            {
                int idx = k * (NY + 1) * (NX + 1) + j * (NX + 1) + i;
                if (isPointExist[idx])
                {
                    fprintf(file, "%f %f %f\n", points[idx].x, points[idx].y, points[idx].z);
                    points_index[idx] = point_num++;
                }
            }
        }
    }
    fprintf(file, "\n");

    // Count active cells and write cell definitions (6 tetrahedrons per active cell)
    int act_num = 0;
    for (int i = 0; i < NX * NY * NZ; i++)
    {
        if (acts[i])
        {
            act_num++;
        }
    }
    fprintf(file, "CELLS %d %d\n", act_num * 6, act_num * 30);

    // Process each active cell, convert it to tetrahedrons, and write their indices
    for (int k = 0; k < NZ; k++)
    {
        for (int j = 0; j < NY; j++)
        {
            for (int i = 0; i < NX; i++)
            {
                int cellIndex = k * NY * NX + j * NX + i;
                if (!acts[cellIndex]) continue;

                int points_index_tmp[8];  // Temporary storage for cell corner indices
                Point<T> points_coords_tmp[8];  // Temporary storage for cell corner coordinates
                int tet_par[24];  // Output tetrahedron indices

                // Collect indices and coordinates of the cell's 8 corners
                for (int index = 0; index < 8; index++)
                {
                    int x = i + (index & 1);
                    int y = j + ((index & 2) >> 1);
                    int z = k + ((index & 4) >> 2);
                    int idx = z * (NY + 1) * (NX + 1) + y * (NX + 1) + x;
                    points_index_tmp[index] = points_index[idx];
                    points_coords_tmp[index] = points[idx];
                }

                // Convert hexahedron to tetrahedrons and write their indices to the file
                hexahedron26tetrahedron(points_index_tmp, points_coords_tmp, tet_par);
                for (int t = 0; t < 6; t++) // Each hexahedron generates 6 tetrahedrons
                {
                    fprintf(file, "4 %d %d %d %d\n", tet_par[t * 4], tet_par[t * 4 + 1], tet_par[t * 4 + 2], tet_par[t * 4 + 3]);
                }
            }
        }
    }
    fprintf(file, "\n");

    // Write cell types (10 indicates tetrahedrons in VTK)
    fprintf(file, "CELL_TYPES %d\n", act_num * 6);
    for (int i = 0; i < act_num * 6; i++)
    {
        fprintf(file, "10\n");
    }
    fprintf(file, "\n");

    // Write fault scalar data for each point
    fprintf(file, "POINT_DATA %d\n", point_num);
    fprintf(file, "SCALARS fault_scalars float 1\n");
    fprintf(file, "LOOKUP_TABLE fault_table\n");
    for (int k = 0; k < NZ + 1; k++)
    {
        for (int j = 0; j < NY + 1; j++)
        {
            for (int i = 0; i < NX + 1; i++)
            {
                int idx = k * (NY + 1) * (NX + 1) + j * (NX + 1) + i;
                if (!isPointExist[idx]) continue;
                int xyidx = j * (NX + 1) + i;
                if (faults[xyidx])
                {
                    fprintf(file, "1\n");
                }
                else
                {
                    fprintf(file, "0\n");
                }
            }
        }
    }
}

int main(int argc, char* argv[])
{
    // Set the number of threads for OpenMP parallelization to 32.
    omp_set_num_threads(32);

    // Define the filename for the grid data file.
    std::string filename = "../../../../../data/Grid_E.out";
    
    // Variables to store the grid dimensions along the x, y, and z axes.
    int NX, NY, NZ;

    // Initialize a variable to accumulate the total execution time.
    double total_time = 0;

    // Start measuring time for reading the grid size.
    auto start = omp_get_wtime();

    // Read the grid dimensions from the file. This function is assumed to populate NX, NY, and NZ.
    readCornerPointGridSize<float>(filename, NX, NY, NZ);

    // Print the dimensions of the grid.
    std::cout << "Grid size: " << NX << " x " << NY << " x " << NZ << std::endl;

    // Allocate vectors for coordinates, corner points, and activity indicators:
    // - `coords` holds x and y coordinates for the grid points.
    // - `zcorns` holds z-coordinates for each cell corner.
    // - `acts` indicates whether each grid cell is active.
    std::vector<float> coords((NX + 1) * (NY + 1) * 2);
    std::vector<float> zcorns(NX * NY * NZ * 8);
    std::vector<bool> acts(NX * NY * NZ);

    // Read the grid data including coordinates, corner points, and active cells.
    readCornerPointGrid<float>(filename, NX, NY, NZ, coords, zcorns, acts);

    // End measuring time and add it to the total execution time.
    auto end = omp_get_wtime();
    total_time += end - start;
    std::cout << "Read grid time: " << end - start << "s" << std::endl;

    // Create vectors to hold grid points and a boolean mask indicating point existence:
    // - `points` stores the 3D coordinates of grid points.
    // - `isPointExist` tracks whether each point is part of the grid.
    std::vector<Point<float>> points((NX + 1) * (NY + 1) * (NZ + 1));
    std::vector<bool> isPointExist((NX + 1) * (NY + 1) * (NZ + 1), false);

    // Start measuring time for assembling the grid points.
    start = omp_get_wtime();

    // Process the grid data to generate points and populate `isPointExist`.
    assemblePoints<float>(coords, zcorns, acts, NX, NY, NZ, points, isPointExist);

    // End measuring time and add it to the total execution time.
    end = omp_get_wtime();
    total_time += end - start;
    std::cout << "Assemble points time: " << end - start << "s" << std::endl;

    // Create a vector to hold fault indicators:
    // - `faults` is a 2D array flattened to 1D, storing fault data for the grid.
    std::vector<bool> faults((NX + 1) * (NY + 1), false);

    // Start measuring time for detecting faults.
    start = omp_get_wtime();

    // Analyze the grid data to identify faults and populate the `faults` vector.
    findFault<float>(coords, isPointExist, NX, NY, NZ, faults);

    // End measuring time and add it to the total execution time.
    end = omp_get_wtime();
    total_time += end - start;
    std::cout << "Find fault time: " << end - start << "s" << std::endl;

    // Start measuring time for writing the grid data to a VTK file.
    start = omp_get_wtime();

    // Write the grid points, cell data, and fault data to a VTK file for visualization.
    print2vtk<float>("../../../../../data/Grid_E.vtk", points, acts, NX, NY, NZ, isPointExist, faults);

    // End measuring time and add it to the total execution time.
    end = omp_get_wtime();
    total_time += end - start;
    std::cout << "Print to vtk time: " << end - start << "s" << std::endl;

    // Print the total execution time for all operations.
    std::cout << "Total time: " << total_time << "s" << std::endl;

    // Return 0 to indicate successful execution.
    return 0;
}
